import tensorflow as tf
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import MaxPooling2D
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Dropout
from tensorflow.keras import Model


class Conv_model(Model):
    def __init__(self):
        # 父类初始化
        super(Conv_model, self).__init__()

        # 添加卷积层
        self.conv_1_1 = Conv2D(input_shape=(160, 160, 3), filters=32, kernel_size=3, padding='same',
                               activation='relu', name='conv_1_1')
        self.conv_1_2 = Conv2D(filters=32, kernel_size=3, padding='same',
                               activation='relu', name='conv_1_2')
        self.conv_1_3 = Conv2D(filters=32, kernel_size=3, padding='same',
                               activation='relu', name='conv_1_3')
        self.conv_1_4 = Conv2D(filters=32, kernel_size=3, padding='same',
                               activation='relu', name='conv_1_4')
        # 添加池化层
        self.max_pool_1 = MaxPooling2D(pool_size=3, name='max_pool_1')

        # dropout
        self.dropout = Dropout(rate=0.5)

        # 继续添加卷积层
        self.conv_2_1 = Conv2D(filters=96, kernel_size=5, padding='same',
                               activation='relu', name='conv_2_1')
        self.conv_2_2 = Conv2D(filters=96, kernel_size=5, padding='same',
                               activation='relu', name='conv_2_2')
        self.conv_2_3 = Conv2D(filters=96, kernel_size=5, padding='same',
                               activation='relu', name='conv_2_3')

        # 池化层
        self.max_pool_2 = MaxPooling2D(pool_size=3, name='max_pool_2')
        self.conv_3_1 = Conv2D(filters=32, kernel_size=3, padding='same',
                               activation='relu', name='conv_2_1')
        self.conv_3_2 = Conv2D(filters=32, kernel_size=3, padding='same',
                               activation='relu', name='conv_2_2')
        self.conv_3_3 = Conv2D(filters=32, kernel_size=3, padding='same',
                               activation='relu', name='conv_2_3')

        # 展平层
        self.flatten = Flatten(name='flatten')

        # 全连接层
        self.dense = Dense(units=10, name='dense', activation='softmax')

    def call(self, x):
        x = self.conv_1_1(x)
        x = self.conv_1_2(x)
        x = self.conv_1_3(x)
        # x = self.conv_1_4(x)
        x = self.max_pool_1(x)
        # x = self.dropout(x)
        x = self.conv_2_1(x)
        x = self.conv_2_2(x)
        x = self.conv_2_3(x)
        x = self.max_pool_2(x)
        # x = self.conv_3_1(x)
        # x = self.conv_3_2(x)
        # x = self.conv_3_3(x)
        x = self.flatten(x)
        x = self.dense(x)
        return x